//
// Created by zhangcc on 2020/7/4.
//

#ifndef OSNET_REID_H
#define OSNET_REID_H

#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <opencv2/opencv.hpp>

#include "cudaWrapper.h"
#include "ioHelper.h"

#define MAX_WORKSPACE_SIZE ((1ULL << 30) * 5) // 1 GB
#define MAX_BINDING_NUMS   (2)

class OsnetReid
{
public:
    OsnetReid(int device_id, int max_batch_size, uint8_t dataType, cv::Size size,
            bool use_dynamic=false, bool alway_serialize=false, bool preprocess=false):
            device_id(device_id), max_batch_size(max_batch_size), dataType(dataType), use_dynamic(use_dynamic), alway_serialize(alway_serialize), preprocess(preprocess)
            {
                inputW = size.width;
                inputH = size.height;
            };

    int init(const std::string &model_file);
    std::vector<std::vector<float>> extract(std::vector<cv::Mat> &inputs);
    ~OsnetReid();

private:
    std::unique_ptr<nvinfer1::ICudaEngine, nvinfer1::Destroy<nvinfer1::ICudaEngine>> mEngine{nullptr};
    std::unique_ptr<nvinfer1::IExecutionContext, nvinfer1::Destroy<nvinfer1::IExecutionContext>> context{nullptr};

    void* bindings[MAX_BINDING_NUMS]{0};
    cudaStream_t stream_t;

    cv::Mat input;

    int device_id = 0;
    int max_batch_size = 3;
    int channels = 3;
    int inputH, inputW;
    int feats_dim = 512;
    bool use_dynamic, alway_serialize;
    uint8_t dataType;
    bool preprocess;
};

#endif //OSNET_REID_H
